#!/usr/bin/env python
# $Id: temp.py 042aa7c43992 2012/12/08 07:49:46 pgurumur $
# Copyright (c) 2012 Prabhu Gurumurthy <pgurumur@gmail.com>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 042aa7c43992 $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-12-08 07:49:46Z $"
# udpqotd - UDP message server

import datetime
import optparse
import os
import SocketServer
import socket
import struct
import sys
import math

from venya.lib import exception
from venya.lib import cmdparser
from venya.lib import ip
from venya.lib import stdlib

PORTNO = 9995

GetTime = lambda value: datetime.date.fromtimestamp(value).strftime(
      "%A %b %d %I:%M:%S %p %Y")

templateDict = {
      1: 'IN_BYTES',
      2: 'IN_PKTS',
      3: 'FLOWS',
      4: 'PROTOCOL',
      5: 'SRC_TOS',
      6: 'TCP_FLAGS',
      7: 'L4_SRC_PORT',
      8: 'IPV4_SRC_ADDR',
      9: 'SRC_MASK',
      10: 'INPUT_SNMP',
      11: 'L4_DST_PORT',
      12: 'IPV4_DST_ADDR',
      13: 'DST_MASK',
      14: 'OUTPUT_SNMP',
      15: 'IPV4_NEXT_HOP',
      16: 'SRC_AS',
      17: 'DST_AS',
      18: 'BGP_IPV4_NEXT_HOP',
      19: 'MUL_DST_PKTS',
      20: 'MUL_DST_BYTES',
      21: 'LAST_SWITCHED',
      22: 'FIRST_SWITCHED',
      23: 'OUT_BYTES',
      24: 'OUT_PKTS',
      25: 'MIN_PKT_LENGTH',
      26: 'MAX_PKT_LENGTH',
      27: 'IPV6_SRC_ADDR',
      28: 'IPV6_DST_ADDR',
      29: 'IPV6_SRC_MASK',
      27: 'IPV6_DST_MASK',
      31: 'IPV6_FLOW_LABEL',
      32: 'ICMP_TYPE',
      33: 'MUL_IGMP_TYPE',
      34: 'SAMPLING_INTERVAL',
      35: 'SAMPLING_ALGORITHM',
      36: 'FLOW_ACTIVE_TIMEOUT',
      37: 'FLOW_INACTIVE_TIMEOUT',
      38: 'ENGINE_TYPE',
      39: 'ENGINE_ID',
      40: 'TOTAL_BYTES_EXP',
      41: 'TOTAL_PKTS_EXP',
      42: 'TOTAL_FLOWS_EXP',
      44: 'IPV4_SRC_PREFIX',
      45: 'IPV4_DST_PREFIX',
      46: 'MPLS_TOP_LABEL_TYPE',
      47: 'MPLS_TOP_LABEL_IPADDR',
      48: 'FLOW_SAMPLER_ID',
      49: 'FLOW_SAMPLER_MODE',
      50: 'FLOW_SAMPLER_RANDOM_INTERVAL',
      52: 'MIN_TTL',
      53: 'MAX_TTL',
      54: 'IPV4_IDENT',
      55: 'DST_TOS',
      56: 'IN_SRC_MAC',
      57: 'OUT_DST_MAC',
      58: 'SRC_VLAN',
      59: 'DST_VLAN',
      60: 'IP_VERSION',
      61: 'DIRECTION',
      62: 'IPV6_NEXT_HOP',
      63: 'BGP_IPV6_NEXT_HOP',
      64: 'IPV6_OPTION_HEADERS',
      70: 'MPLS_LABEL1',
      71: 'MPLS_LABEL2',
      72: 'MPLS_LABEL3',
      73: 'MPLS_LABEL4',
      74: 'MPLS_LABEL5',
      75: 'MPLS_LABEL6',
      76: 'MPLS_LABEL7',
      77: 'MPLS_LABEL8',
      78: 'MPLS_LABEL9',
      79: 'MPLS_LABEL10',
      80: 'IN_DST_MAC',
      81: 'OUT_SRC_MAC',
      82: 'IF_NAME',
      83: 'IF_DESC',
      84: 'SAMPLER_NAME',
      85: 'IN_PERM_BYTES',
      86: 'IN_PERM_PKTS',
      88: 'FRAGMENT_OFFSET',
      89: 'FORWARDING_STATUS',
      90: 'MPLS_PAL_RD',
      91: 'MPLS_PREFIX_LEN',
      92: 'SRC_TRAFFIC_INDEX',
      93: 'DST_TRAFFIC_INDEX',
      94: 'APP_DESCR',
      95: 'APP_TAG',
      96: 'APP_NAME',
      98: 'DIFFSERV_PT',
      99: 'REPLICATION_FACTOR',
      102: 'L2_PKT_SECTION_OFFSET',
      103: 'L2_PKT_SECTION_SIZE',
      104: 'L2_PKT_SECTION_DATA'
}

class NetFlowPacketHeader(object):
   def __init__(self, HeaderBuffer):
      self._ver = 0
      self._count = 0
      self._uptime = 0
      self._unixtm = 0
      self._seq = 0
      self._srcid = 0

      if len(HeaderBuffer) == 20:
         temp = struct.unpack(">HHIIII", HeaderBuffer)
         self._ver = temp[0]
         self._count = temp[1]
         self._uptime = temp[2]
         self._unixtm = temp[3]
         self._seq = temp[4]
         self._srcid = temp[5]
      else:
         self.error("Invalid netflow packet header size, length: %d", len(
            HeaderBuffer))

      print(self._srcid)

   def error(self, fmt, *args):
      mesg = stdlib.sprintf(fmt, *args)
      raise ValueError(mesg)

   def __str__(self):
      return str(self._srcid)

   def __int__(self):
      return int(self._count)

   version = property(fget = lambda self: self._ver)
   sequence = property(fget = lambda self: self._seq)
   sourceid = property(fget = lambda self: self._srcid)
   count = property(fget = lambda self: self._count)

class TemplateHeader(object):
   def __init__(self, Buffer):
      super(TemplateHeader, self).__init__()
      self._dict = {
            1: 'IN_BYTES',
            2: 'IN_PKTS',
            3: 'FLOWS',
            4: 'PROTOCOL',
            5: 'SRC_TOS',
            6: 'TCP_FLAGS',
            7: 'L4_SRC_PORT',
            8: 'IPV4_SRC_ADDR',
            9: 'SRC_MASK',
            10: 'INPUT_SNMP',
            11: 'L4_DST_PORT',
            12: 'IPV4_DST_ADDR',
            13: 'DST_MASK',
            14: 'OUTPUT_SNMP',
            15: 'IPV4_NEXT_HOP',
            16: 'SRC_AS',
            17: 'DST_AS',
            18: 'BGP_IPV4_NEXT_HOP',
            19: 'MUL_DST_PKTS',
            20: 'MUL_DST_BYTES',
            21: 'LAST_SWITCHED',
            22: 'FIRST_SWITCHED',
            23: 'OUT_BYTES',
            24: 'OUT_PKTS',
            25: 'MIN_PKT_LENGTH',
            26: 'MAX_PKT_LENGTH',
            27: 'IPV6_SRC_ADDR',
            28: 'IPV6_DST_ADDR',
            29: 'IPV6_SRC_MASK',
            27: 'IPV6_DST_MASK',
            31: 'IPV6_FLOW_LABEL',
            32: 'ICMP_TYPE',
            33: 'MUL_IGMP_TYPE',
            34: 'SAMPLING_INTERVAL',
            35: 'SAMPLING_ALGORITHM',
            36: 'FLOW_ACTIVE_TIMEOUT',
            37: 'FLOW_INACTIVE_TIMEOUT',
            38: 'ENGINE_TYPE',
            39: 'ENGINE_ID',
            40: 'TOTAL_BYTES_EXP',
            41: 'TOTAL_PKTS_EXP',
            42: 'TOTAL_FLOWS_EXP',
            44: 'IPV4_SRC_PREFIX',
            45: 'IPV4_DST_PREFIX',
            46: 'MPLS_TOP_LABEL_TYPE',
            47: 'MPLS_TOP_LABEL_IPADDR',
            48: 'FLOW_SAMPLER_ID',
            49: 'FLOW_SAMPLER_MODE',
            50: 'FLOW_SAMPLER_RANDOM_INTERVAL',
            52: 'MIN_TTL',
            53: 'MAX_TTL',
            54: 'IPV4_IDENT',
            55: 'DST_TOS',
            56: 'IN_SRC_MAC',
            57: 'OUT_DST_MAC',
            58: 'SRC_VLAN',
            59: 'DST_VLAN',
            60: 'IP_VERSION',
            61: 'DIRECTION',
            62: 'IPV6_NEXT_HOP',
            63: 'BGP_IPV6_NEXT_HOP',
            64: 'IPV6_OPTION_HEADERS',
            70: 'MPLS_LABEL1',
            71: 'MPLS_LABEL2',
            72: 'MPLS_LABEL3',
            73: 'MPLS_LABEL4',
            74: 'MPLS_LABEL5',
            75: 'MPLS_LABEL6',
            76: 'MPLS_LABEL7',
            77: 'MPLS_LABEL8',
            78: 'MPLS_LABEL9',
            79: 'MPLS_LABEL10',
            80: 'IN_DST_MAC',
            81: 'OUT_SRC_MAC',
            82: 'IF_NAME',
            83: 'IF_DESC',
            84: 'SAMPLER_NAME',
            85: 'IN_PERM_BYTES',
            86: 'IN_PERM_PKTS',
            88: 'FRAGMENT_OFFSET',
            89: 'FORWARDING_STATUS',
            90: 'MPLS_PAL_RD',
            91: 'MPLS_PREFIX_LEN',
            92: 'SRC_TRAFFIC_INDEX',
            93: 'DST_TRAFFIC_INDEX',
            94: 'APP_DESCR',
            95: 'APP_TAG',
            96: 'APP_NAME',
            98: 'DIFFSERV_PT',
            99: 'REPLICATION_FACTOR',
            102: 'L2_PKT_SECTION_OFFSET',
            103: 'L2_PKT_SECTION_SIZE',
            104: 'L2_PKT_SECTION_DATA'
      }

      val = math.floor(len(Buffer) / 2)
      format = None
      if len(Buffer) % 2:
         format = ">%dHb" %int(val)
      else:
         format = ">%dH" %int(val)

      temp = struct.unpack(format, Buffer)
      # print(temp)

class FlowHeader(object):
   def __init__(self, TemplateBuffer, Buffer):
      temp = struct.unpack("!HHHH", TemplateBuffer)
      # print("flowset id: %d, length: %d bytes, template id: %d, count: %d" %(
         # temp[0], temp[1], temp[2], temp[3]))
         #
      print("flowset id: %d, length: %d" %(temp[0], temp[1]))
      print("template id: %d, field count: %d" %(temp[2], temp[3]))
      if temp[0] == 1:
         print("Template Record with options")
      elif temp[0] == 0:
         print("Template Record")
         theader = TemplateHeader(Buffer)
      elif temp[0] > 255:
         """
         val = math.floor(len(Buffer) / 4)
         if len(Buffer) % 4:
            format = ">2H%dIb" %(int(val) - 1)
         else:
            format = ">2H%dI" %(int(val) - 1)

         foo = struct.unpack(format, Buffer)
         """
         # print(len(Buffer))
         # print("Data Record")

      # print(temp)

class handler(SocketServer.DatagramRequestHandler):
   def handle(self):
      packet = ""
      for ix in self.rfile.readlines():
         packet += ix

      pktlen = len(packet)
      if pktlen == 0:
         print("ignoring netflow header, because the size is zero!")
      elif pktlen < 20:
         version = struct.unpack("!H", packet[:2])
         print("invalid netflow header size, length: %d, version: %d" %(
            pktlen, version[0]))
      elif pktlen >= 20:
         pkthdr = NetFlowPacketHeader(packet[:20])
         print("pkt length: %d, version: %d, record count: %d, sequence: %d" %(
            pktlen, pkthdr.version, pkthdr.count, pkthdr.sequence))
         flowtype = struct.unpack("!HHHH", packet[20:28])
         if flowtype[0] == 0:
            print("template flowset, length: %d, template id: %d, count: %d" %(
               flowtype[1], flowtype[2], flowtype[3]))
            mod = pktlen % 2
            if mod != 0:
               pktlen -= 1

            count = (pktlen - 28) / 2
            hformat = "!%dH" %count
            records = struct.unpack(hformat, packet[28:pktlen])
            count = 0
            for value in records:
               value = int(value)
               if (count % 2 == 0) and (value in templateDict):
                  tempdict[templateDict[value]] = None

               count += 1

            print(count)
         elif flowtype[0] == 1:
            print("options template flowset, length: %d, template id: %d" %(
               flowtype[1], flowtype[2]))
            optlength = struct.unpack("!H", packet[28:30])
            print("option scope length: %d, option length: %d" %(
               flowtype[3], optlength[0]))
            records = struct.unpack("!HH", packet[30:34])
            print(records)
         elif flowtype[0] > 255:
            print("data flowset, id: %d, length: %d" %(
               flowtype[0], flowtype[1]))

            mod = pktlen % 4
            if mod != 0:
               pktlen -= mod

            count = (pktlen - 28) / 4
            dformat = "!%dI" %count
            records = struct.unpack(dformat, packet[28:pktlen])
            print(records)

      # print int(NetFlowPacketHeader(newmsg[:20]))
      # FlowHeader(newmsg[20:28], newmsg[28:])

      """
      if len(newmsg) > 20:
         strlen = len(newmsg) - 40
         # format = ">hhIIII%ds" %strlen
         format = ">hhIIIIIIIII%ds" %strlen
         temp = struct.unpack(format, newmsg)
        # print "Client %s said ``%s''" % (self.client_address[0], newmsg)
        # self.wfile.write(self.server.oldmsg)
      self.server.oldmsg = newmsg
      """


"""
s = SocketServer.UDPServer(('',PORTNO), handler)
print "Awaiting UDP messages on port %d" % PORTNO
s.oldmsg = "This is the starting message."
s.serve_forever()
"""

class Server(object):
   def __init__(self, Address = '', Port = 0):
      self._addr = None
      self._port = 0

      if Address:
         self._addr = Address

      if Port:
         self._port = int(Port)

      self._server = SocketServer.UDPServer(('', self._port), handler)

   def listen(self):
      self._server.serve_forever()

class netflowd(cmdparser.GetOpt):
   option_list = (
         optparse.make_option("-p", "--port", type = "int",
            help = "UDP port to listen, default: 9995", default = 9995,
            dest = "port"),
   ) + cmdparser.GetOpt.option_list

   def __init__(self, **Keywords):
      super(netflowd, self).__init__(**Keywords)
      self._server = None

   def handle(self):
      (options, args) = self.parse()
      if options.port:
         self._server = Server('', options.port)
         self._server.listen()

if __name__ == "__main__":
   flowd = None
   try:
      flowd = netflowd(version = 1.0, args = sys.argv)
      if len(flowd) <= 1:
         flowd.printhelp()
      else:
         flowd.handle()
   except exception.GetOptError as err:
      print(err)
      sys.exit(1)

